//
// Created by 98770 on 2023/6/22.
//

#include "serverControl.h"

static uint16_t IDmap[512] = {0};
int capacity = 0;
static int backupCounter = 0;
char *DNSlist[100];
#if _WIN32
extern HANDLE cacheLock; //保护dynamicCache
extern HANDLE tableLock; //保护id映射表
#elif __linux__
extern pthread_mutex_t cacheLock;
extern pthread_mutex_t tableLock;
#endif

bool initStaticCache(struct hashmap_s *hashmap, char *filename) {
    //从文件中加载静态缓存
    if (!createHashmap(hashmap)) {
        printf("create hashmap failed when init static cache!\n");
        return false;
    }
    if (!loadCacheFromDisk(hashmap, filename)) {
        printf("load static cache failed!\n");
        return false;
    }
    return true;
}

bool initDynamicCache(struct hashmap_s *hashmap) {
    //从文件中加载动态缓存
    if (!createHashmap(hashmap)) {
        printf("create hashmap failed when init dynamic cache!\n");
        return false;
    }
    return true;
}

bool initServer(struct hashmap_s *staticHashmap, char *filename, struct hashmap_s *dynamicHashmap) {
    if (!initStaticCache(staticHashmap, filename))
        return false;
    if (!initDynamicCache(dynamicHashmap))
        return false;
    DNSListIn("..//ISPserver.txt");
//    initMutex();
    printf("init server successfully!\n");
    printf("Num of entries in static cache: %d\n", hashmap_num_entries(staticHashmap));
    return true;
}

char *resolveByStaticCache(struct hashmap_s *hashmap, char *domainName) {
    //从静态缓存中解析
    cacheNode *node = getRecordFromHashmap(hashmap, domainName);
    if (node == NULL)
        return "NULL";
    else {
        //如果记录已经过期，,但是还未来得及则删除
        if (node->ttl <= 0) {
            deleteRecordFromHashmap(hashmap, domainName);
            return "NULL";
        }
            //如果记录未过期，则返回
        else {
            return node->ip;
        }
    }
}

char *resolveByDynamicCache(struct hashmap_s *hashmap, char *domainName) {
    //从动态缓存中解析
#if _WIN32
    WaitForSingleObject(cacheLock, INFINITE);
#elif __linux__
//    pthread_mutex_lock(&cacheLock);
#endif
    cacheNode *node = getRecordFromHashmap(hashmap, domainName);
    if (node == NULL)
        return "NULL";
    else {
        //如果记录已经过期，,但是还未来得及则删除
        if (node->ttl <= 0) {
            deleteRecordFromHashmap(hashmap, domainName);
#if _WIN32
            ReleaseMutex(cacheLock);
#elif __linux__
//            pthread_mutex_unlock(&cacheLock);
#endif
            return "NULL";
        }
            //如果记录未过期，则返回
        else {
#if _WIN32
            ReleaseMutex(cacheLock);
#elif __linux__
//            pthread_mutex_unlock(&cacheLock);
#endif
            return node->ip;
        }
    }
}

char *resolveQuery(struct hashmap_s *staticHashmap, struct hashmap_s *dynamicHashmap, char *packet, int len,
                   char *DNSserver) {
    char *domainName = getDomainName(packet, len);
    cacheNode *node;
    //按顺序从静态缓存、动态缓存、互联网解析
    char *ip = resolveByStaticCache(staticHashmap, domainName);
    if (strcmp(ip, "NULL") == 0) {
        ip = resolveByDynamicCache(dynamicHashmap, domainName);
        if (strcmp(ip, "NULL") == 0) {
            node = resolveByInternet(packet, len, DNSserver);
            ip = node->ip;
            if (strcmp(ip, "NULL") == 0) {
                free(domainName);
                return "NULL";
            } else {
                putRecordInHashmap(dynamicHashmap, domainName, ip, node->ttl);
                free(domainName);
                return ip;
            }
        } else {
            free(domainName);
            return ip;
        }
    } else {
        free(domainName);
        return ip;
    }
}

cacheNode *resolveByInternet(char *packet, int len, char *DNSserver) {
    //从指定的DNS服务器解析
    char *addr = malloc(sizeof(char) * 16);
    struct DNS_answer *ans = resolveQueryOnISP(packet, len, DNSserver);
    cacheNode *node = malloc(sizeof(cacheNode));
    if (ans == NULL) {
        return NULL;
    } else {
        strcpy(node->ip, ans->ip);
        strcpy(node->domain, ans->domain);
        node->ttl = ans->ttl;
        free(ans);
        return node;
    }
}

//从DNS请求报文中获取域名
char *getDomainName(char *packet, int len) {

    const char *position = packet + sizeof(struct dns_header);
    char *domain_name = NULL;
    int domain_name_length = 0;
    int max_domain_name_length = 255;

    domain_name = (char *) malloc(max_domain_name_length + 1); // 分配内存，加1是为了容纳字符串结尾的'\0'
    if (domain_name == NULL) {
        printf("Failed to allocate memory for domain_name\n");
        return NULL;
    }

    // 解析报文开始位置
    while (*position != '\0') {
        // 判断是否为指针
        if ((*position & 0xC0) == 0xC0) {
            // 是指针，解析指针并跳过
            int pointer_offset = ntohs(*(uint16_t *) position) & 0x3FFF;
            position = packet + pointer_offset;
        } else {
            // 非指针，提取标签内容并添加到域名
            int label_length = *position;

            // 检查域名长度是否超过最大限制
            if (domain_name_length + label_length + 1 > max_domain_name_length) {
                printf("Domain name exceeds the maximum length\n");
                free(domain_name);
                return NULL;
            }

            // 拷贝标签内容到域名
            strncpy(domain_name + domain_name_length, position + 1, label_length);
            domain_name_length += label_length;
            domain_name[domain_name_length] = '.';
            domain_name_length++;

            position += label_length + 1;
        }
    }

    // 移除末尾多余的点号
    if (domain_name_length > 0 && domain_name[domain_name_length - 1] == '.') {
        domain_name[domain_name_length - 1] = '\0';
        domain_name_length--;
    }

    // 重新调整内存大小以适应实际域名长度
    char *resized_domain_name = (char *) realloc(domain_name, domain_name_length + 1);
    if (resized_domain_name != NULL) {
        domain_name = resized_domain_name;
    }

    return domain_name;
}

int buildResponse(char *request, char *response, char *ip, int len, uint32_t ttl) {
    memcpy(response, request, len);
    char *position = response + len;
    struct dns_header *header = (struct dns_header *) response;
    //当ip为0.0.0.0，直接返回域名不存在
    if (strcmp(ip, "0.0.0.0") == 0) {
        header->flags = htons(0x8183);
        header->qdcount = htons(1);
        header->ancount = htons(0);
        header->nscount = htons(0);
        header->arcount = htons(0);
        return len;
    }
    header->flags = htons(0x8180);
    header->qdcount = htons(1);
    header->ancount = htons(1);
    header->nscount = htons(0);
    header->arcount = htons(0);
    struct DNSResourceRecord *resourceRecord = (struct DNSResourceRecord *) position;
    resourceRecord->name = htons(0xc00c);
    resourceRecord->type = htons(0x0001);
    resourceRecord->class = htons(0x0001);
    resourceRecord->ttl = htonl(ttl);
    resourceRecord->rdlength = htons(0x004);
    resourceRecord->rdata = inet_addr(ip);
    printf("Response IP: %s\n", ip);
    return len + sizeof(struct DNSResourceRecord);
}

bool processQuery(int s, struct sockaddr_in *si_in, struct hashmap_s *staticHashmap, struct hashmap_s *dynamicHashmap,
                  char *packet, int len, char *DNSserver) {

    char *ip;
    uint32_t ttl;
    int slen = sizeof(*si_in);
    int responseLen;
    static char outBuffer[1024];
    static char senderIp[16];
    strcpy(senderIp, inet_ntoa(si_in->sin_addr));
    IDAssignForQuery(packet);
    char *domainName = getDomainName(packet, len);
    printf("Domain name to be resolved: %s\n", domainName);
    ip = resolveByStaticCache(staticHashmap, domainName);
    if (strcmp(ip, "NULL") != 0) {
        printf("Resolved IP: %s on static list\n", ip);
        buildResponse(packet, outBuffer, ip, len, 0xFFFFFFFF);
        IDAssignForResponse(outBuffer);
        sendto(s, outBuffer, 1024, 0, (struct sockaddr *) si_in, slen);
        free(domainName);
        return true;
    }
    ip = resolveByDynamicCache(dynamicHashmap, domainName);
    if (strcmp(ip, "NULL") != 0) {
        printf("Resolved IP: %s on dynamic list\n", ip);
        ttl = resolveTTLonCache(staticHashmap, domainName);
        buildResponse(packet, outBuffer, ip, len, ttl);
        IDAssignForResponse(outBuffer);
        sendto(s, outBuffer, len, 0, (struct sockaddr *) si_in, slen);
        free(domainName);
        return true;
    }
    printf("Find in ISP DNS server\n");
    struct DNS_answer *response = resolveQueryOnISP(packet, len, DNSserver);
    if (response == NULL || strcpy(response->ip, "0.0.0.0") == 0) {
        //回应为域名不存在，或为配制的静态解析列表中0.0.0.0的域名
        ttl = resolveTTLonCache(staticHashmap, domainName);
        buildResponse(packet, outBuffer, "0.0.0.0", len, ttl);
        IDAssignForResponse(outBuffer);
        sendto(s, outBuffer, len, 0, (struct sockaddr *) si_in, slen);
        free(domainName);
        return true;
    }
    storeTypeARecord(dynamicHashmap, response->packet, len);
    memcpy(outBuffer, response->packet, response->len);
    IDAssignForResponse(outBuffer);
    sendto(s, outBuffer, response->len, 0, (struct sockaddr *) si_in, slen);
    free(response);
    return true;
}

void *refreshCacheForUnix(void *cache) {
    cache = (struct hashmap_s *) cache;
    printf("reach to cache refresh call \n");
    while (true) {
        sleep(1);
        backupCounter++;
        if(backupCounter==60){
            backupCounter=0;
            storeCacheToDisk(cache);
        }
        refreshHashmap(cache);
    }
}

void refreshCacheForWin32(void *cache) {
    cache = (struct hashmap_s *) cache;
    printf("reach to cache refresh call \n");
    while (true) {
        sleep(1);
        backupCounter++;
        if(backupCounter==60){
            backupCounter=0;
#if _WIN32
            WaitForSingleObject(cacheLock, 1);
            storeCacheToDisk(cache);
            ReleaseMutex(cacheLock);
#endif
        }
#if _WIN32
        WaitForSingleObject(cacheLock, 1);
        refreshHashmap(cache);
        ReleaseMutex(cacheLock);
#endif
    }
}

uint32_t resolveTTLonCache(struct hashmap_s *hashmap, char *domainName) {
    cacheNode *node = getRecordFromHashmap(hashmap, domainName);
    if (node == NULL) {
        return 0;
    }
    return node->ttl;
}

void storeTypeARecord(struct hashmap_s *hashmap, char *response, int len) {
    struct dns_header *header = (struct dns_header *) response;
    uint16_t answerCount = ntohs(header->ancount);
    //从回应中提取出所有的资源记录，如果记录类型为A，就加入到缓存中
    //reader指向第一个资源记录的开头
    char *reader = response + len;
    for (int i = 0; i < answerCount; i++) {
        struct DNSResourceRecord *resourceRecord = (struct DNSResourceRecord *) reader;
        char *domainName;
        if ((resourceRecord->name & 0x00ff) != 0xc0) {
            domainName= malloc(256);
            strcpy(domainName, getDomainName(response, len));
        }
        domainName = getFullName(response, resourceRecord->name);
        if(domainName==NULL)
            break;
        if (ntohs(resourceRecord->type) == 1) {
            char *ip = inet_ntoa(*(struct in_addr *) &resourceRecord->rdata);
#if _WIN32
            WaitForSingleObject(cacheLock, INFINITE);
#elif __linux__
//            pthread_mutex_lock(&cacheLock);
#endif
            putRecordInHashmap(hashmap, domainName, ip, ntohl(resourceRecord->ttl));
#if _WIN32
            ReleaseMutex(cacheLock);
#elif __linux__
//            pthread_mutex_unlock(&cacheLock);
#endif
        }
        reader += sizeof(struct DNSResourceRecord) + ntohs(resourceRecord->rdlength) - sizeof(uint32_t);
        if (domainName != NULL)
            free(domainName);
    }
}

char *getFullName(char *dnsHeader, uint16_t offset) {
    char *fullName = NULL;
    int fullNameIndex = 0;
    int jumped = 0;
    int nameEnd = 0;
    int maxFullNameLength = 256; // 假设最大长度为256

    fullName = (char *) malloc(maxFullNameLength);
    if (fullName == NULL) {
        printf("Failed to allocate memory for fullName\n");
        return NULL;
    }
    offset = ntohs(offset)&0x3ff;
    const char *reader = dnsHeader + offset;

    while (*reader != 0 && !jumped && !nameEnd) {
        if ((*reader & 0xC0) == 0xC0) {
            int jumpOffset = ntohs(*(uint16_t *) reader) & 0x3FFF;
            reader = dnsHeader + jumpOffset;
            jumped = 1;
        } else {
            int labelLength = *reader;

            if ((labelLength & 0xC0) == 0xC0) {
                int jumpOffset = ntohs(*(uint16_t *) reader) & 0x3FFF;
                const char *jumpReader = dnsHeader + jumpOffset;
                int jumpLength = *(jumpReader++);
                int currentPos = fullNameIndex;

                while (jumpLength > 0 && fullNameIndex < maxFullNameLength - 1) {
                    fullName[fullNameIndex++] = jumpReader[currentPos];
                    currentPos++;
                    jumpLength--;
                }

                fullName[fullNameIndex++] = '.';
                jumped = 1;
            } else {
                reader++;

                for (int i = 0; i < labelLength && fullNameIndex < maxFullNameLength - 1; i++) {
                    fullName[fullNameIndex++] = *reader++;
                }

                fullName[fullNameIndex++] = '.';

                if (labelLength == 0) {
                    nameEnd = 1;
                }
            }
        }
    }

    fullName[fullNameIndex - 1] = '\0';

    return fullName;
}

int DNSListIn(char *filename) {
    FILE *fp;
    char *line = NULL;
    size_t len = 0;
    char ip[16];
    fp = fopen(filename, "r");
    if (fp == NULL) {
        printf("open file failed\n");
        DNSlist[0] = malloc(16);
        strcpy(DNSlist[0], DEFAULT_DNS_SERVER);
        return 1;
    }
    capacity = 0;
    while (fscanf(fp, "%s\n", ip) != EOF && capacity < 100) {
        DNSlist[capacity] = malloc(16);
        strcpy(DNSlist[capacity], ip);
        capacity++;
    }
    fclose(fp);
    return capacity;
}

char *DNSGet() {
    static int index = 0;
    index++;
    return DNSlist[index % capacity];
}

void DNSFree() {
    for (int i = 0; i < capacity; i++) {
        if (DNSlist[i] != NULL)
            free(DNSlist[i]);
    }
}

void initIDmap(struct hashmap_s *hashmap) {
    const unsigned initial_size = 20;
    if (0 != hashmap_create(initial_size, hashmap)) {
        // error!
    }
}

int16_t findAnUsableID() {
    //generate random id
    srand(time(NULL));
    uint16_t num = rand();
    num = num % 512;
    while (IDmap[num] != 0) {
        num = rand();
        num = num % 512;
    }
    return num;
};

void IDAssignForQuery(char *packet) {
#if _WIN32
    WaitForSingleObject(tableLock, INFINITE);
#elif __linux__
//    pthread_mutex_lock(&tableLock);
#endif
    uint16_t id = findAnUsableID();
    uint16_t origin = *(uint16_t *) (packet + 0);
    IDmap[id] = origin;
    *(uint16_t *) (packet + 0) = id;
#if _WIN32
    ReleaseMutex(tableLock);
#elif __linux__
//    pthread_mutex_unlock(&tableLock);
#endif
}

void IDAssignForResponse(char *packet) {
#if _WIN32
    WaitForSingleObject(tableLock, INFINITE);
#elif __linux__
//    pthread_mutex_lock(&tableLock);
#endif
    uint16_t id = *(uint16_t *) (packet + 0);
    uint16_t origin = IDmap[id];
    *(uint16_t *) (packet + 0) = origin;
    IDmap[id] = 0;
#if _WIN32
    ReleaseMutex(tableLock);
#elif __linux__
//    pthread_mutex_unlock(&tableLock);
#endif
}